import onnx_graphsurgeon as gs
import onnx
import argparse
import sys

def main():
    parser = argparse.ArgumentParser(description="裁剪 ONNX 模型，删除指定起始和结束节点之间的子图")
    parser.add_argument("-i", "--input", required=True, help="输入的 ONNX 文件路径")
    parser.add_argument("-o", "--output", required=True, help="输出的 ONNX 文件路径")
    parser.add_argument("-s", "--start", required=True, help="起始节点名称")
    parser.add_argument("-e", "--end", required=True, help="结束节点名称")
    
    args = parser.parse_args()
    
    try:
        # 加载 ONNX 模型
        model = onnx.load(args.input)
        graph = gs.import_onnx(model)
        
        # 检查节点是否存在
        tensors = graph.tensors()
        if args.start not in tensors:
            print(f"错误：找不到起始节点 '{args.start}'")
            print(f"可用的节点：{list(tensors.keys())}")
            sys.exit(1)
        
        if args.end not in tensors:
            print(f"错误：找不到结束节点 '{args.end}'")
            print(f"可用的节点：{list(tensors.keys())}")
            sys.exit(1)
        
        # 裁剪从起始节点到结束节点的子图
        start_tensor = tensors[args.start]
        end_tensor = tensors[args.end]
        
        # 使用图遍历找到从起始节点到结束节点的所有节点
        def find_connected_nodes(start_tensor, end_tensor):
            # 从起始张量开始，找到所有可达的节点
            visited_nodes = []
            visited_tensor_names = set()
            queue = [start_tensor]
            visited_tensor_names.add(start_tensor.name)
            
            while queue:
                current_tensor = queue.pop(0)
                
                # 找到使用当前张量作为输入的节点
                for node in graph.nodes:
                    if current_tensor in node.inputs and node not in visited_nodes:
                        visited_nodes.append(node)
                        
                        # 将节点的输出张量加入队列
                        for output_tensor in node.outputs:
                            if output_tensor.name not in visited_tensor_names:
                                visited_tensor_names.add(output_tensor.name)
                                queue.append(output_tensor)
                                
                                # 如果到达目标张量，停止搜索
                                if output_tensor.name == end_tensor.name:
                                    return visited_nodes
            
            return visited_nodes
        
        # 找到连接起始和结束节点的所有节点
        connected_nodes = find_connected_nodes(start_tensor, end_tensor)
        
        if not connected_nodes:
            print(f"错误：无法找到从 '{args.start}' 到 '{args.end}' 的连接路径")
            sys.exit(1)
        
        # 移除不在连接路径上的节点
        nodes_to_remove = []
        for node in graph.nodes:
            if node not in connected_nodes:
                nodes_to_remove.append(node)
        
        for node in nodes_to_remove:
            graph.nodes.remove(node)
        
        # 设置输入和输出
        graph.inputs = [start_tensor]
        graph.outputs = [end_tensor]
        
        # 清理未使用的张量和节点
        graph.cleanup()
        
        # 保存裁剪后的模型
        onnx.save(gs.export_onnx(graph), args.output)
        
        print(f"成功裁剪模型：{args.input} -> {args.output}")
        
    except FileNotFoundError:
        print(f"错误：找不到文件 {args.input}")
        sys.exit(1)
    except Exception as e:
        print(f"错误：{str(e)}")
        sys.exit(1)

if __name__ == "__main__":
    main()